#include "table.h"

#include "mem.h"
#include "assert.h"

#include <limits.h>

#define T Table_T

struct T {
    int size;
    int (*cmp)(const void *x, const void *y);
    unsigned int (*hash)(const void *key);
    int length;
    unsigned int timestamp;

    struct binding {
        struct binding *link;
        const void *key;
        void *value;
    } **buckets;
};

static int
cmpatom(const void *x, const void *y)
{
    return x != y;
}

static unsigned int
hashatom(const void *key)
{
    return (size_t)key >> 2; // address aligning, the right 2 bit may be 0
}

T Table_new(int hint, int cmp(const void *x, const void *y), unsigned int hash(const void *key))
{
    static int primes[] = {509, 509, 1021, 2053, 4093, 8191, 16381, 32771, 65221, INT_MAX};
    int i, size;

    assert(hint >= 0);
    for (i = 1; primes[i] <= hint && primes[i] <= INT_MAX; i++)
        ;
    size = primes[i - 1];

    T table = (T)ALLOC(sizeof(*table) + size * sizeof(table->buckets[0]));
    table->size = size;
    table->cmp = cmp != NULL ? cmp : cmpatom;
    table->hash = hash != NULL ? hash : hashatom;
    table->length = 0;
    table->timestamp = 0;
    table->buckets = (struct binding**)(table + 1);

    for (i = 0; i < table->size; i++)
        table->buckets[i] = NULL;

    return table;
}

void Table_free(T *table)
{
    struct binding *p, *q;

    assert(table && *table);

    if ((*table)->length > 0) {
        for (int i = 0; i < (*table)->size; i++) {
            for (p = (*table)->buckets[i]; p; p = q) {
                q = p->link;
                FREE(p);
            }
        }
    }

    FREE(*table);
}

int Table_length(T table)
{
    assert(table);
    return table->length;
}

static struct binding*
search_key(T table, const void *key)
{
    struct binding *p;
    int slot = table->hash(key) % table->size;
    for (p = table->buckets[slot]; p; p = p->link) {
        if (table->cmp(p->key, key) == 0)
            return p;
    }
    return NULL;
}

void* Table_put(T table, const void *key, void *value)
{
    struct binding *p;
    void *prev = NULL;

    assert(table);
    assert(key);

    p = search_key(table, key);
    if (p == NULL) {
        int slot = table->hash(key) % table->size;

        NEW(p);
        p->key = key;
        p->link = table->buckets[slot];
        table->buckets[slot] = p;
        table->length++;

    } else {
        prev = p->value;
    }
    
    p->value = value;
    table->timestamp++;

    return prev;
}

void* Table_get(T table, const void *key)
{
    struct binding *p; 

    assert(table);
    assert(key);

    p = search_key(table, key);
    return p ? p->value : NULL;
}

void* Table_remove(T table, const void *key)
{
    struct binding **pp;
    int slot = table->hash(key) % table->size;
    void *value = NULL;

    assert(table);
    assert(key);
    table->timestamp++;

    for (pp = &table->buckets[slot]; *pp; pp = &(*pp)->link) {
        if (table->cmp((*pp)->key, key) == 0) {
            struct binding *p = *pp;
            value = p->value;
            *pp = p->link;
            FREE(p);
            table->length--;
            break;
        }
    }
    return value;
}

void Table_map(T table, void apply(const void *key, void **value, void *cl), void *cl)
{
    struct binding *p;
    unsigned int stamp = table->timestamp;

    assert(table);
    assert(apply);

    for (int i = 0; i < table->size; i++) {
        for (p = table->buckets[i]; p; p = p->link) {
            apply(p->key, &p->value, cl);
            assert(table->timestamp == stamp);
        }
    }
}

void** Table_toArray(T table, void *end)
{
    struct binding *p;
    void **array;

    assert(table);

    array = ALLOC((2 * table->length + 1) * sizeof(*array));

    int j = 0;
    for (int i = 0; i < table->size; i++) {
        for (p = table->buckets[i]; p; p = p->link) {
            array[j++] = (void*)p->key;
            array[j++] = p->value;
        }
    }
    array[j] = end;

    return array;
}

